#include "bzbnet/tls.h"
#include "bzbstring/stringutils.h"
#include "bzbio/io.h"

#include <gnutls/gnutls.h>
#include <gnutls/x509.h>

#include <memory>
#include <cstring>


gnutls_certificate_credentials_t cred = nullptr;

bool tls_initialized = false;


void throw_gnutls_error(int err) {

	throw std::runtime_error(
		bzbstring::concat("TLS Error: GnuTLS: ",
							gnutls_strerror(err)));

}


void bzbnet::initialize_tls() {

	gnutls_global_init();

	cred = nullptr;
	int ret = gnutls_certificate_allocate_credentials(&cred);
	if(ret != GNUTLS_E_SUCCESS) {
		throw_gnutls_error(ret);
	}

	ret = gnutls_certificate_set_x509_system_trust(cred);
	if(ret == 0) {
		throw std::runtime_error("TLS Error: No certificates found");
	}

	if(ret < 0) {
		throw_gnutls_error(ret);
	}

	tls_initialized = true;
}


void bzbnet::destroy_tls() {

	gnutls_certificate_free_credentials(cred);
	tls_initialized = false;
}


bzbnet::TLSClient::TLSClient() {

	socket = TCPSocket();

	if(!tls_initialized) {
		initialize_tls();
	}
}


bzbnet::TLSClient::~TLSClient() {

	socket.close();
}


void bzbnet::TLSClient::connect(std::string hostname, std::string service) {

	if(!tls_initialized) {
		initialize_tls();
	}

	// Initialize TLS session
	gnutls_session_t session;
	int ret = gnutls_init(&session, GNUTLS_CLIENT);
	if(ret != GNUTLS_E_SUCCESS) {
		throw_gnutls_error(ret);
	}

	// Set default priority
	ret = gnutls_set_default_priority(session);
	if(ret != GNUTLS_E_SUCCESS) {
		throw_gnutls_error(ret);
	}

	// Set trusted certificates
	ret = gnutls_credentials_set(session, GNUTLS_CRD_CERTIFICATE, cred);
	if(ret != GNUTLS_E_SUCCESS) {
		throw_gnutls_error(ret);
	}

	socket.connect(hostname, service);

	// Associate the session to the socket's file descriptor
	gnutls_transport_set_ptr(session,
		(gnutls_transport_ptr_t) (uintptr_t) socket.get_raw_fd());

	ret = gnutls_server_name_set(session, GNUTLS_NAME_DNS,
								 hostname.c_str(), hostname.size());
	if(ret != GNUTLS_E_SUCCESS) {
		throw_gnutls_error(ret);
	}

	// Perform handshake
	ret = gnutls_handshake(session);
	if(ret != GNUTLS_E_SUCCESS) {
		throw_gnutls_error(ret);
	}

	unsigned certs_size = 0;
	const gnutls_datum_t* certs = gnutls_certificate_get_peers(session, &certs_size);
	if(certs == nullptr || certs_size == 0) {
		throw std::runtime_error("TLS Error: Couldn't retrieve peer certificate");
	}

	unsigned status = (unsigned) -1;
	ret = gnutls_certificate_verify_peers2(session, &status);
	if(ret != GNUTLS_E_SUCCESS) {
		throw_gnutls_error(ret);
	}

	if(status != 0) {

		gnutls_datum_t msg;

#if GNUTLS_VERSION_NUMBER >= 0x030104

		gnutls_certificate_type_t type = gnutls_certificate_type_get(session);
		ret = gnutls_certificate_verification_status_print(status, type, &msg, 0);

#else
		ret = -1;
#endif

		if(ret == 0) {
		
			gnutls_free(msg.data);
			throw std::runtime_error(bzbstring::concat("TLS Error: GnuTLS: ", msg.data));

		} else {

			throw std::runtime_error(bzbstring::concat(
				"TLS Error: GnuTLS: Certificate validation failed with status: ", status));
		}
	}

#if GNUTLS_VERSION_NUMBER >= 0x030104

	ret = gnutls_certificate_verify_peers3(session, hostname.c_str(), &status);
	if(ret) {
		throw_gnutls_error(ret);
	}

#else

	gnutls_x509_crt_t cert;
	ret = gnutls_x509_crt_init(&cert);
	if(ret) {
		throw_gnutls_error(ret);
	}

	ret = gnutls_x509_crt_import(cert, certs, GNUTLS_X509_FMT_DER);
	if(ret) {
		throw_gnutls_error(ret);
	}

	ret = gnutls_x509_crt_check_hostname(cert, host);
	if(ret == 0) {
		throw std::runtime_error("TLS Error: Hostname does not match certificate");
	}

	gnutls_x509_crt_deinit(cert);

#endif

	gnutls_session_t* buff = new gnutls_session_t();
	memcpy((void*) buff, (void*) &session, sizeof(gnutls_session_t));
	session_ptr = (void*) buff;
}


std::vector<char> bzbnet::TLSClient::recv(unsigned int packetsize) {

	if(!session_ptr) {
		throw std::runtime_error("TLS Error: Cannot receive data prior to connection");
	}

	std::vector<char> data;
	char* packet = new char[packetsize];
	unsigned int totalsize = 0;
	int size = 0;

	do {

		size = gnutls_record_recv(*((gnutls_session_t*) session_ptr),
								 packet, packetsize);

		bzbio::println(packet);

		if(size < 0) {
			throw_gnutls_error(size);
		}

		if(size) {
			data.insert(data.begin() + totalsize, packet, packet + size);
			totalsize += size;
		}

	} while(size > 0);

	if(!totalsize) {

		// Connection closed
		gnutls_deinit(*((gnutls_session_t*) session_ptr));
		delete (gnutls_session_t*) session_ptr;
		session_ptr = nullptr;
		socket.close();
	}

	return data;
}


void bzbnet::TLSClient::send(std::string data) {

	if(!session_ptr) {
		throw std::runtime_error("TLS Error: Cannot send data prior to connection");
	}

	int ret = gnutls_record_send(*((gnutls_session_t*) session_ptr),
								 &data.at(0), data.size());
	if(ret < 0) {
		throw_gnutls_error(ret);
	}

}


void bzbnet::TLSClient::send(std::vector<char> data) {

	if(!session_ptr) {
		throw std::runtime_error("TLS Error: Cannot send data prior to connection");
	}

	int ret = gnutls_record_send(*((gnutls_session_t*) session_ptr),
								 &data.at(0), data.size());
	if(ret < 0) {
		throw_gnutls_error(ret);
	}

}


void bzbnet::TLSClient::close() {

	if(session_ptr) {
		int ret = gnutls_bye(*((gnutls_session_t*) session_ptr), GNUTLS_SHUT_RDWR);
		if(ret < 0) {
			throw_gnutls_error(ret);
		}

		gnutls_deinit(*((gnutls_session_t*) session_ptr));
		delete (gnutls_session_t*) session_ptr;

		socket.close();
	}

}


bool bzbnet::TLSClient::is_open() {

	return session_ptr != nullptr;
}


int bzbnet::TLSClient::get_raw_fd() {

	return socket.get_raw_fd();
}

